# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
__author__: xyh
__date__: 2016-09-13
用于将各种公告的基础消息发送到RabbitMQ, 供后端取出来，将pdf转换为html
各种公告包括：A股预披露, A股公告, 新三板预披露, 新三版公告
可以通过announcementTime, 或某份公告文件的_id来选择信息发送

通过在当前文件路径下的config.cfg设置发送的queue_name
设置config.cfg的[publish]下的某一个time_range来设置公告发布的时间区间, 例如'pre_xsb_time_range' = 2016-05-06 ~ 2016-09-12
"""

import sys
import json
import traceback
import requests
import pymongo
import pika
import time
from xtls.timeparser import now
from scpy.rabbit_asy_producer import AsyPublisher
from scpy.logger import get_logger

reload(sys)
sys.setdefaultencoding('utf-8')
logger = get_logger(__file__)

QINIU = {
    'CHECK_URL': 'http://10.132.23.104:7300/api/check',
    'HOST': '10.132.23.104',
    'PORT': 7300,
    'BUCKET': 'sc-crawler',
}

# --------------------Base Class of BuildMessage-------------------------------
class BuildMessage(object):
    """
    基类：构建公告信息的message, 为producer准备message
    """
    MONGO = pymongo.MongoClient('10.132.23.104', 27017)

    def __init__(self, start_anno_time, stop_anno_time, STARTUP_TIME, _id=''):
        self.db = ''
        self.collection = ''
        self._id = _id
        self.STARTUP_TIME = STARTUP_TIME
        self.start_anno_time = start_anno_time
        self.stop_anno_time = stop_anno_time
        self.error_count = 0

    def check_exist_in_qiniu(self, anno):
        file_key = anno.get('fileKey')
        check = requests.get(QINIU['CHECK_URL'], {'filename': file_key, 'bucket': QINIU['BUCKET']})
        if check.status_code / 100 != 2:
            return False
        return True

    def check_upload_error(self, anno):
        file_key = anno.get('fileKey')
        upload_error = self.MONGO['qiniu']['uploadError'].find_one({'key': file_key})
        if upload_error:
            self.error_count += 1

    def update_publish_time(self, data):
        self.MONGO[self.db][self.collection].update_one({'_id': data.get('_id')}, {'$set': {'publishTime': now()}})

    def find_from_mongo(self):
        """
        从mongo中取出符合规则的公告文件的元数据, 可以根据需要重写该方法
        :return: cursor,可以当生成器来用, 符合规则的公告文件的元数据
        """
        if self._id:
            return self.MONGO[self.db][self.collection].find({'_id': self._id,
                                                              'fileKey': {'$ne': None},
                                                              '$or': [{'publishTime': {'$lt': self.STARTUP_TIME}},
                                                                      {'publishTime': None}]
                                                              })

        return self.MONGO[self.db][self.collection].find(f)

    def build_message(self, data):
        """
        由子类去实现,通过从db取出的data, 生成message
        :param data:
        :return: dict message
        """
        raise NotImplementedError

    def completed(self):
        all_count = self.find_from_mongo().count()
        if all_count == 0:
            return True
        if all_count == self.error_count:
            return True

        return False

    def run(self):
        for anno in self.find_from_mongo():
            if not anno.get('fileKey', None):
                logger.info('Non fileKey {}'.format(anno.get('_id', '')))
                continue

            if not self.check_exist_in_qiniu(anno):
                self.check_upload_error(anno)
                continue

            message = self.build_message(anno)

            yield message

            self.update_publish_time(anno)


# ---------------------A股公告的BuildMessage-------------------------------------
class BuildShenhu(BuildMessage):
    def __init__(self, start_anno_time, stop_anno_time, STARTUP_TIME, _id=''):
        super(BuildShenhu, self).__init__(start_anno_time, stop_anno_time, STARTUP_TIME, _id)
        self.db = 'stockProj'
        self.collection = 'shenhuStockAnno'
        self.name = 'shenhu'

    def build_message(self, data):
        message = {
            'announcementTime': data.get('announcementTime') or '',
            'columnId': data.get('columnId') or '',
            'title': data.get('title') or '',
            'fileKey': data.get('fileKey') or '',
            'orgId': data.get('orgId') or '',
            'stockCode': data.get('stockCode') or '',
            'type': data.get('type') or [],
        }
        return message


# ---------------------新三板公告的BuildMessage-------------------------------------
class BuildXsb(BuildMessage):
    def __init__(self, start_anno_time, stop_anno_time, STARTUP_TIME, _id=''):
        super(BuildXsb, self).__init__(start_anno_time, stop_anno_time, STARTUP_TIME, _id)
        self.db = 'stockProj'
        self.collection = 'xinsanban'
        self.name = 'xsb'

    def build_message(self, data):
        message = {
            'announcementTime': data.get('announcementTime') or '',
            'columnId': data.get('columnId') or '',
            'title': data.get('title') or '',
            'fileKey': data.get('fileKey') or '',
            'orgId': data.get('orgId') or '',
            'stockCode': data.get('stockCode') or '',
            'type': data.get('type') or [],
        }
        return message


# ---------------------A股预披露的BuildMessage-------------------------------------
class BuildPreAShare(BuildMessage):
    def __init__(self, start_anno_time, stop_anno_time, STARTUP_TIME, _id=''):
        super(BuildPreAShare, self).__init__(start_anno_time, stop_anno_time, STARTUP_TIME, _id)
        self.db = 'stockProj'
        self.collection = 'preDisclosure'
        self.name = 'pre_shenhu'

    def build_message(self, data):
        message = {
            'announcementTime': data.get('announcementTime') or '',
            'columnId': data.get('columnId') or '',
            'title': data.get('title') or '',
            'fileKey': data.get('fileKey') or '',
            'orgId': data.get('orgId') or '',
            'stockCode': data.get('stockCode') or '',
            'type': data.get('type') or [],
            'stockMarket': 'sh',
        }
        return message

    def find_from_mongo(self):
        """
        从mongo中取出符合规则的公告文件的元数据, 在重写基类方法，增加了一个筛选条件'stockMarket': 'sh'
        :return: cursor,可以当生成器来用, 符合规则的公告文件的元数据
        """

        if self._id:
            return self.MONGO[self.db][self.collection].find({'_id': self._id,
                                                              'fileKey': {'$ne': None},
                                                              'stockMarket': 'sh',
                                                              '$or': [{'publishTime': {'$lt': self.STARTUP_TIME}},
                                                                      {'publishTime': None}]
                                                              })

        return self.MONGO[self.db][self.collection].find({'announcementTime': {'$gte': self.start_anno_time, '$lte': self.stop_anno_time},
                                                          'fileKey': {'$ne': None},
                                                          'stockMarket': 'sh',
                                                          '$or': [{'publishTime': {'$lt': self.STARTUP_TIME}},
                                                                  {'publishTime': None}]
                                                          })


# ---------------------新三板预披露的BuildMessage-------------------------------------
class BuildPreXsb(BuildMessage):
    def __init__(self, start_anno_time, stop_anno_time, STARTUP_TIME, _id=''):
        super(BuildPreXsb, self).__init__(start_anno_time, stop_anno_time, STARTUP_TIME, _id)
        self.db = 'stockProj'
        self.collection = 'preDisclosure'
        self.name = 'pre_xsb'

    def build_message(self, data):
        message = {
            'announcementTime': data.get('announcementTime') or '',
            'columnId': data.get('columnId') or '',
            'title': data.get('title') or '',
            'fileKey': data.get('fileKey') or '',
            'orgId': data.get('orgId') or '',
            'stockCode': data.get('stockCode') or '',
            'type': data.get('type') or [],
            'stockMarket': 'xsb',
        }
        return message

    def find_from_mongo(self):
        """
        从mongo中取出符合规则的公告文件的元数据, 在重写基类方法，增加了一个筛选条件'stockMarket': 'xsb'
        :return: cursor,可以当生成器来用, 符合规则的公告文件的元数据
        """
        if self._id:
            return self.MONGO[self.db][self.collection].find({'_id': self._id,
                                                              'fileKey': {'$ne': None},
                                                              'stockMarket': 'xsb',
                                                              '$or': [{'publishTime': {'$lt': self.STARTUP_TIME}},
                                                                      {'publishTime': None}]
                                                              })

        return self.MONGO[self.db][self.collection].find({'announcementTime': {'$gte': self.start_anno_time, '$lte': self.stop_anno_time},
                                                          'fileKey': {'$ne': None},
                                                          'stockMarket': 'xsb',
                                                          '$or': [{'publishTime': {'$lt': self.STARTUP_TIME}},
                                                                  {'publishTime': None}]
                                                          })


# ---------------------------producer-------------------------
class MongoAsyProducer(AsyPublisher):
    def __init__(self, amqp_url, queue_name, build_message):
        super(MongoAsyProducer, self).__init__(amqp_url, queue_name)
        self.build_message = build_message
        self.PUBLISH_INTERVAL = 1800  # 间隔30分钟又publish一次

    def publish_message(self):
        message_generator = self.build_message.run()
        while True:
            try:
                message = message_generator.next()
            except StopIteration:
                break

            properties = pika.BasicProperties(app_id='example-publisher',
                                              content_type='application/json',
                                              headers=message)
            self._channel.basic_publish('',
                                        self.QUEUE,
                                        json.dumps(message),
                                        properties)

            logger.info('publish done %s: %s' % (self.build_message.name, message.get('fileKey')))

        logger.info('[x] %s: Completed one time of messages publish' % self.build_message.name)
        time.sleep(60)

        if self.build_message.completed():
            self.stop()
        else:
            logger.info('%s: Continue Next Publish_Message, please Wait %d seconds' % (self.build_message.name, self.PUBLISH_INTERVAL))
            logger.info('to_publish:%d, error:%d' % (self.build_message.find_from_mongo().count(), self.build_message.error_count))
            self.schedule_next_message()


def produce_message(amqp_url, queue_name, build_message):
    producer = MongoAsyProducer(amqp_url=amqp_url,
                                queue_name=queue_name,
                                build_message=build_message)
    try:
        producer.run()
    except KeyboardInterrupt:
        producer.stop()
    except Exception, e:
        trace = traceback.format_exc()
        logger.exception(trace)
        producer.stop()


